import argparse

from rocket.builder.container_builder import ComponentContainerBuilder
from rocket.builder.resolver import ReferenceResolver
from rocket.builder.resource_loader import YamlResourceLoader, JsonResourceLoader, PyResourceLoader
from rocket.estimator.estimator import Estimator


def parse_args():
    argparser = argparse.ArgumentParser()
    argparser.add_argument("config", type=str)
    argparser.add_argument("-m", type=str, default="train")
    argparser.add_argument("-e", type=int, default=1)

    args = argparser.parse_args()
    return args


def get_resource_loader(config):
    if config.endswith(".yml") or config.endswith(".yaml"):
        return YamlResourceLoader()
    elif config.endswith(".json"):
        return JsonResourceLoader()
    elif config.endswith(".py"):
        return PyResourceLoader()
    else:
        raise NotImplemented()


if __name__ == '__main__':
    args = parse_args()
    loader = get_resource_loader(args.config)
    resolver = ReferenceResolver()
    builder = ComponentContainerBuilder(loader, resolver)

    container = builder.build(args.config)
    estimator = Estimator(container)

    if args.m == "train":
        estimator.train(epochs=args.e, context="cuda")
    elif args.m == "val":
        estimator.val(epochs=args.e, context="cuda")
